# Copyright 2025 HuggingFace Inc., THUDM, and the LlamaFactory team.
#
# This code is inspired by the HuggingFace's transformers library and the THUDM's ChatGLM implementation.
# https://github.com/huggingface/transformers/blob/v4.40.0/examples/pytorch/summarization/run_summarization.py
# https://github.com/THUDM/ChatGLM-6B/blob/main/ptuning/main.py
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional

import numpy as np
import torch
from transformers.utils import is_jieba_available, is_nltk_available

from ...extras.constants import IGNORE_INDEX
from ...extras.misc import numpify
from ...extras.packages import is_rouge_available


if TYPE_CHECKING:
    from transformers import EvalPrediction, PreTrainedTokenizer


if is_jieba_available():
    import jieba  # type: ignore


if is_nltk_available():
    from nltk.translate.bleu_score import SmoothingFunction, sentence_bleu  # type: ignore


if is_rouge_available():
    from rouge_chinese import Rouge  # type: ignore

"""
在评估时，如果直接保存完整的 logits（形状为 (batch_size, seq_len, vocab_size)），会占用大量内存。例如：
batch=8，seq_len=2048，vocab_size=32000 → 单个 logits 张量约 1.98G 内存（float32）4*8*2048*32000/1024/1024/1024=1.953125GB
因此，在只需要预测 token ID 的场景下（比如计算 accuracy），提前用 argmax 处理可以大幅节省内存。
"""

# 定义 eval_logit_processor 函数，用于评估时处理模型输出 logits，减少内存占用
def eval_logit_processor(logits: "torch.Tensor", labels: "torch.Tensor") -> "torch.Tensor":
    r"""Compute the token with the largest likelihood to reduce memory footprint."""
    # 如果 logits 是列表或元组，则表示可能包含额外信息（如 MoE 的 aux loss）
    if isinstance(logits, (list, tuple)):
        # 如果第一个元素是三维张量（batch_size, seq_len, vocab_size），则认为是主 logits
        if logits[0].dim() == 3:  # (batch_size, seq_len, vocab_size)
            logits = logits[0]
        else:  # moe models have aux loss
            # 对于 MoE 模型，可能存在辅助损失，此时取第二个元素作为主 logits
            logits = logits[1]

    # 检查 logits 是否为三维张量，否则抛出异常
    if logits.dim() != 3:
        raise ValueError("Cannot process the logits.")

    # 使用 torch.argmax 获取每个位置上概率最大的 token ID，用于计算准确率等指标
    return torch.argmax(logits, dim=-1)


# 使用 dataclass 定义 ComputeAccuracy 类，用于计算准确率并支持 batch_eval_metrics
@dataclass
class ComputeAccuracy:
    r"""Compute accuracy and support `batch_eval_metrics`."""

    # _dump 方法用于汇总当前累积的指标结果，并重置缓存，返回 float 类型的平均值字典
    def _dump(self) -> Optional[dict[str, float]]:
        result = None
        # 如果存在 score_dict 属性，则对其中每个 key 的值取平均并转换为 float
        if hasattr(self, "score_dict"):
            result = {k: float(np.mean(v)) for k, v in self.score_dict.items()}

        # 重置 score_dict 缓存，准备下一轮评估
        self.score_dict = {"accuracy": []}
        return result

    # __post_init__ 是 dataclass 的特殊方法，在初始化后自动调用
    def __post_init__(self):
        # 初始化时清空或创建 score_dict
        self._dump()

    # 实现 Callable 接口，使得该类可以像函数一样被调用，用于处理 EvalPrediction 并计算准确率
    def __call__(self, eval_preds: "EvalPrediction", compute_result: bool = True) -> Optional[dict[str, float]]:
        # 将 predictions 和 label_ids 转换为 numpy 格式以便后续处理
        preds, labels = numpify(eval_preds.predictions), numpify(eval_preds.label_ids)
        # 遍历每个样本进行准确率计算
        for i in range(len(preds)):
            # 去除预测中的最后一个 token，标签从第二个 token 开始（对应 shift right）
            pred, label = preds[i, :-1], labels[i, 1:]
            # 创建 label_mask，忽略 IGNORE_INDEX 对应位置, 就是-100
            label_mask = label != IGNORE_INDEX
            # 计算当前样本中预测正确的 token 占比，并添加到 score_dict 中
            self.score_dict["accuracy"].append(np.mean(pred[label_mask] == label[label_mask]))

        # 如果需要返回最终结果，则调用 _dump 返回当前批次的平均准确率
        if compute_result:
            return self._dump()


# 使用 dataclass 定义 ComputeSimilarity 类，用于评估生成文本与标签之间的相似度
@dataclass
class ComputeSimilarity:
    r"""Compute text similarity scores and support `batch_eval_metrics`.

    Wraps the tokenizer into metric functions, used in CustomSeq2SeqTrainer.
    """

    # 初始化时需要传入一个 tokenizer，用于解码预测结果和标签
    tokenizer: "PreTrainedTokenizer"

    # _dump 方法用于汇总当前累积的指标结果，并重置缓存，返回 float 类型的平均值字典
    def _dump(self) -> Optional[dict[str, float]]:
        result = None
        # 如果存在 score_dict 属性，则对其中每个 key 的值取平均并转换为 float
        if hasattr(self, "score_dict"):
            result = {k: float(np.mean(v)) for k, v in self.score_dict.items()}

        # 重置 score_dict 缓存，准备下一轮评估，初始化包含 rouge 和 bleu 指标
        self.score_dict = {"rouge-1": [], "rouge-2": [], "rouge-l": [], "bleu-4": []}
        return result

    # __post_init__ 是 dataclass 的特殊方法，在初始化后自动调用
    def __post_init__(self):
        # 初始化时清空或创建 score_dict
        self._dump()

    # 实现 Callable 接口，使得该类可以像函数一样被调用，用于处理 EvalPrediction 并计算相似度
    def __call__(self, eval_preds: "EvalPrediction", compute_result: bool = True) -> Optional[dict[str, float]]:
        # 将 predictions 和 label_ids 转换为 numpy 格式以便后续处理
        preds, labels = numpify(eval_preds.predictions), numpify(eval_preds.label_ids)

        # 将 IGNORE_INDEX 替换为 pad_token_id，防止解码出错
        preds = np.where(preds != IGNORE_INDEX, preds, self.tokenizer.pad_token_id)
        labels = np.where(labels != IGNORE_INDEX, labels, self.tokenizer.pad_token_id)

        # 使用 tokenizer 批量解码预测结果和真实标签
        decoded_preds = self.tokenizer.batch_decode(preds, skip_special_tokens=True)
        decoded_labels = self.tokenizer.batch_decode(labels, skip_special_tokens=True)

        # 遍历每一对解码后的预测和标签，进行分词并计算 ROUGE 和 BLEU 分数
        for pred, label in zip(decoded_preds, decoded_labels):
            # 对预测和标签分别使用 jieba 进行中文分词
            hypothesis = list(jieba.cut(pred))
            reference = list(jieba.cut(label))

            # 如果预测或参考为空，则 ROUGE 分数设为 0.0
            if len(" ".join(hypothesis).split()) == 0 or len(" ".join(reference).split()) == 0:
                result = {"rouge-1": {"f": 0.0}, "rouge-2": {"f": 0.0}, "rouge-l": {"f": 0.0}}
            else:
                # 使用 Rouge 库计算 ROUGE-1、ROUGE-2、ROUGE-L 分数
                rouge = Rouge()
                scores = rouge.get_scores(" ".join(hypothesis), " ".join(reference))
                result = scores[0]

            # 将每个 ROUGE 分数的 F1 值加入到 score_dict 中
            for k, v in result.items():
                self.score_dict[k].append(round(v["f"] * 100, 4))

            # 计算 BLEU-4 分数，并加入 score_dict
            bleu_score = sentence_bleu([list(label)], list(pred), smoothing_function=SmoothingFunction().method3)
            self.score_dict["bleu-4"].append(round(bleu_score * 100, 4))

        # 如果需要返回最终结果，则调用 _dump 返回当前批次的平均指标
        if compute_result:
            return self._dump()
